{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xoPiGVD8LNma"
   },
   "source": [
    "# Buffer\n",
    "\n",
    "The replay buffer is a fundamental component in deep reinforcement learning (DRL) implementations. In Tianshou, the Buffer module extends the functionality of the Batch class by providing trajectory tracking capabilities and sampling utilities that go beyond basic data storage.\n",
    "\n",
    "Tianshou provides several buffer implementations, with `ReplayBuffer` and `VectorReplayBuffer` being the most fundamental. The latter is specifically designed for parallelized environments, which will be covered in the [Vectorized Environment](https://tianshou.readthedocs.io/en/master/02_notebooks/L3_Vectorized__Environment.html) tutorial. This tutorial focuses exclusively on the `ReplayBuffer` implementation."
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from tianshou.data import Batch, ReplayBuffer"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OdesCAxANehZ"
   },
   "source": [
    "## Core Functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fUbLl9T_SrTR"
   },
   "source": [
    "### Circular Queue Storage Mechanism\n",
    "\n",
    "The buffer stores data in batches using a circular queue mechanism. When the buffer reaches its maximum capacity, newly added data automatically overwrites the oldest entries, ensuring efficient memory utilization while maintaining the most recent experiences."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "mocZ6IqZTH62",
    "outputId": "66cc4181-c51b-4a47-aacf-666b92b7fc52"
   },
   "source": [
    "# Initialize buffer with maximum capacity of 10 transitions\n",
    "print(\"========================================\")\n",
    "dummy_buf = ReplayBuffer(size=10)\n",
    "print(dummy_buf)\n",
    "print(f\"maxsize: {dummy_buf.maxsize}, data length: {len(dummy_buf)}\")\n",
    "\n",
    "# Add 3 transition steps sequentially\n",
    "print(\"========================================\")\n",
    "for i in range(3):\n",
    "    dummy_buf.add(\n",
    "        Batch(obs=i, act=i, rew=i, terminated=0, truncated=0, done=0, obs_next=i + 1, info={}),\n",
    "    )\n",
    "print(dummy_buf)\n",
    "print(f\"maxsize: {dummy_buf.maxsize}, data length: {len(dummy_buf)}\")\n",
    "\n",
    "# Add 10 additional transitions to demonstrate circular queue behavior\n",
    "# Note: First 3 transitions will be overwritten as capacity is exceeded\n",
    "print(\"========================================\")\n",
    "for i in range(3, 13):\n",
    "    dummy_buf.add(\n",
    "        Batch(obs=i, act=i, rew=i, terminated=0, truncated=0, done=0, obs_next=i + 1, info={}),\n",
    "    )\n",
    "print(dummy_buf)\n",
    "print(f\"maxsize: {dummy_buf.maxsize}, data length: {len(dummy_buf)}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "H8B85Y5yUfTy"
   },
   "source": [
    "### Batch-Compatible Operations\n",
    "\n",
    "Consistent with the `Batch` interface, `ReplayBuffer` supports standard operations including concatenation, splitting, advanced slicing, and indexing."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cOX-ADOPNeEK",
    "outputId": "f1a8ec01-b878-419b-f180-bdce3dee73e6"
   },
   "source": [
    "print(dummy_buf[-1])\n",
    "print(dummy_buf[-3:])\n",
    "# Additional Batch methods can be explored as needed"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vqldap-2WQBh"
   },
   "source": [
    "### Persistence and Serialization\n",
    "\n",
    "The buffer can be serialized to disk while preserving trajectory information. This capability is particularly valuable for offline reinforcement learning applications, where pre-collected experience datasets are utilized for training."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "id": "Ppx0L3niNT5K"
   },
   "source": [
    "_dummy_buf = pickle.loads(pickle.dumps(dummy_buf))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Eqezp0OyXn6J"
   },
   "source": [
    "### Reserved Keys for DRL Integration\n",
    "\n",
    "To facilitate seamless integration with DRL algorithms, `ReplayBuffer` utilizes nine reserved keys within the `Batch` structure. These keys follow the [Gymnasium](https://gymnasium.farama.org/index.html#) conventions:\n",
    "\n",
    "*   `obs` - Current observation\n",
    "*   `act` - Action taken\n",
    "*   `rew` - Reward received\n",
    "*   `terminated` - Episode termination flag (goal reached or failure)\n",
    "*   `truncated` - Episode truncation flag (time limit or external interruption)\n",
    "*   `done` - Combined termination/truncation indicator\n",
    "*   `obs_next` - Subsequent observation\n",
    "*   `info` - Auxiliary information dictionary\n",
    "*   `policy` - Policy-specific data\n",
    "\n",
    "**Best Practice**: Use the `info` dictionary for custom metadata rather than adding additional top-level keys. The `done` flag is internally tracked to determine trajectory boundaries, episode lengths, and cumulative rewards.\n",
    "\n",
    "```python\n",
    "# Not recommended: Custom top-level keys\n",
    "buf.add(Batch(......, extra_info=0))\n",
    "\n",
    "# Recommended: Use info dictionary\n",
    "buf.add(Batch(......, info={\"extra_info\": 0}))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ueAbTspsc6jo"
   },
   "source": [
    "### Experience Sampling\n",
    "\n",
    "The primary function of a replay buffer in DRL is to enable experience sampling for training. The buffer provides two methods for this purpose:\n",
    "\n",
    "1. `ReplayBuffer.sample()` - Direct batch sampling with specified size\n",
    "2. `ReplayBuffer.split(..., shuffle=True)` - Split buffer into multiple batches with optional shuffling"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "P5xnYOhrchDl",
    "outputId": "bcd2c970-efa6-43bb-8709-720d38f77bbd"
   },
   "source": [
    "dummy_buf.sample(batch_size=5)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IWyaOSKOcgK4"
   },
   "source": [
    "## Trajectory Management\n",
    "\n",
    "A distinguishing feature of `ReplayBuffer` compared to `Batch` is its trajectory tracking capability, which maintains episode boundaries and associated metadata.\n",
    "\n",
    "The following example demonstrates trajectory tracking by simulating three episodes:\n",
    "1. First episode: 3 steps (completed)\n",
    "2. Second episode: 5 steps (completed)\n",
    "3. Third episode: 5 steps (ongoing)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "editable": true,
    "id": "H0qRb6HLfhLB",
    "outputId": "9bdb7d4e-b6ec-489f-a221-0bddf706d85b",
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "trajectory_buffer = ReplayBuffer(size=10)\n",
    "\n",
    "# Episode 1: 3 steps, terminates at step 2\n",
    "print(\"========================================\")\n",
    "for i in range(3):\n",
    "    result = trajectory_buffer.add(\n",
    "        Batch(\n",
    "            obs=i,\n",
    "            act=i,\n",
    "            rew=i,\n",
    "            terminated=1 if i == 2 else 0,\n",
    "            truncated=0,\n",
    "            done=i == 2,\n",
    "            obs_next=i + 1,\n",
    "            info={},\n",
    "        ),\n",
    "    )\n",
    "    print(result)\n",
    "print(trajectory_buffer)\n",
    "print(f\"maxsize: {trajectory_buffer.maxsize}, data length: {len(trajectory_buffer)}\")\n",
    "\n",
    "# Episode 2: 5 steps, terminates at step 7\n",
    "print(\"========================================\")\n",
    "for i in range(3, 8):\n",
    "    result = trajectory_buffer.add(\n",
    "        Batch(\n",
    "            obs=i,\n",
    "            act=i,\n",
    "            rew=i,\n",
    "            terminated=1 if i == 7 else 0,\n",
    "            truncated=0,\n",
    "            done=i == 7,\n",
    "            obs_next=i + 1,\n",
    "            info={},\n",
    "        ),\n",
    "    )\n",
    "    print(result)\n",
    "print(trajectory_buffer)\n",
    "print(f\"maxsize: {trajectory_buffer.maxsize}, data length: {len(trajectory_buffer)}\")\n",
    "\n",
    "# Episode 3: 5 steps added, episode still ongoing\n",
    "print(\"========================================\")\n",
    "for i in range(8, 13):\n",
    "    result = trajectory_buffer.add(\n",
    "        Batch(obs=i, act=i, rew=i, terminated=0, truncated=0, done=False, obs_next=i + 1, info={}),\n",
    "    )\n",
    "    print(result)\n",
    "print(trajectory_buffer)\n",
    "print(f\"maxsize: {trajectory_buffer.maxsize}, data length: {len(trajectory_buffer)}\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dO7PWdb_hkXA"
   },
   "source": [
    "### Episode Metrics Tracking\n",
    "\n",
    "The `ReplayBuffer.add()` method returns a tuple containing four values: `(current_index, episode_reward, episode_length, episode_start_index)`. \n",
    "\n",
    "**Important**: The `episode_reward` and `episode_length` fields are only populated when an episode completes (i.e., when `done=True`). This automatic computation eliminates the need for manual episode metric tracking during data collection."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xbVc90z8itH0"
   },
   "source": [
    "### Episode Boundary Navigation\n",
    "\n",
    "The buffer provides mechanisms to navigate episode boundaries efficiently. Consider the following scenario where we query a mid-episode step:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4mKwo54MjupY",
    "outputId": "9ae14a7e-908b-44eb-afec-89b45bac5961"
   },
   "source": [
    "print(trajectory_buffer)\n",
    "print(\"========================================\")\n",
    "\n",
    "data = trajectory_buffer[6]\n",
    "print(data)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "p5Co_Fmzj8Sw"
   },
   "source": [
    "### Determining Episode Start Indices\n",
    "\n",
    "Step 6 belongs to the second episode (steps 3-7). While this may appear straightforward, determining the episode start index programmatically is non-trivial due to:\n",
    "\n",
    "1. **Ambiguous done flags**: The preceding `done` flag approach fails when the buffer contains incomplete episodes, as step 3 is surrounded by `done=False` values\n",
    "2. **Complex buffer structures**: Advanced buffers like `VectorReplayBuffer` do not store data sequentially, making boundary detection more challenging\n",
    "\n",
    "The buffer provides a unified API to handle these complexities through the `prev()` method, which identifies the previous step within an episode:"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "source": [
    "# Query previous steps for indices [0, 1, 2, 3, 4, 5, 6]\n",
    "# Episode boundaries prevent backward traversal past episode starts\n",
    "print(trajectory_buffer.prev(np.array([0, 1, 2, 3, 4, 5, 6])))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4Wlb57V4lQyQ"
   },
   "source": [
    "The output confirms that step 3 marks the episode start. The complementary `ReplayBuffer.next()` method enables forward traversal to identify episode terminations, providing a consistent interface across all buffer implementations."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zl5TRMo7oOy5",
    "outputId": "4a11612c-3ee0-4e74-b028-c8759e71fbdb"
   },
   "source": [
    "# Query next steps for indices [4, 5, 6, 7, 8, 9]\n",
    "# Episode boundaries prevent forward traversal past episode ends\n",
    "print(trajectory_buffer.next(np.array([4, 5, 6, 7, 8, 9])))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YJ9CcWZXoOXw"
   },
   "source": [
    "### Identifying Incomplete Episodes\n",
    "\n",
    "The buffer maintains tracking of incomplete episodes through the `unfinished_index()` method, which identifies the most recent step of ongoing episodes (marked with `done=False`):"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Xkawk97NpItg",
    "outputId": "df10b359-c2c7-42ca-e50d-9caee6bccadd"
   },
   "source": [
    "print(trajectory_buffer.unfinished_index())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8_lMr0j3pOmn"
   },
   "source": [
    "### Applications in DRL Algorithms\n",
    "\n",
    "These trajectory navigation APIs are essential for computing algorithmic quantities such as:\n",
    "- Generalized Advantage Estimation (GAE)\n",
    "- N-step returns\n",
    "- Temporal difference targets\n",
    "\n",
    "The unified interface ensures modular design and enables algorithm implementations that generalize across different buffer types. For reference implementations, see the [Tianshou policy base class](https://github.com/thu-ml/tianshou/blob/6fc68578127387522424460790cbcb32a2bd43c4/tianshou/policy/base.py#L384)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FEyE0c7tNfwa"
   },
   "source": [
    "## Advanced Topics\n",
    "\n",
    "### Specialized Buffer Implementations\n",
    "\n",
    "Tianshou provides several specialized buffer variants for advanced use cases:\n",
    "\n",
    "*   **PrioritizedReplayBuffer**: Implements [prioritized experience replay](https://arxiv.org/abs/1511.05952) for importance-weighted sampling\n",
    "*   **CachedReplayBuffer**: Maintains a primary buffer with auxiliary cached buffers for improved sample efficiency in specific scenarios\n",
    "*   **ReplayBufferManager**: Base class for custom buffer implementations requiring management of multiple buffer instances\n",
    "\n",
    "Consult the API documentation and source code for detailed implementation specifications.\n",
    "\n",
    "### Recurrent Neural Network Support\n",
    "\n",
    "The buffer initialization accepts a `stack_num` parameter (default: 1) to enable frame stacking for recurrent neural network (RNN) integration in DRL algorithms. This feature facilitates temporal sequence processing by automatically stacking consecutive observations. Refer to the API documentation for configuration details and usage examples."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
